{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 尝试SAM模型\n",
    "import numpy as np\n",
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "from segment_anything import SamPredictor, sam_model_registry\n",
    "import os\n",
    "os.environ[\"KMP_DUPLICATE_LIB_OK\"]=\"TRUE\"   # 一个粗暴的方法，解决库重复初始化问题\n",
    "\n",
    "\n",
    "\n",
    "sam_checkpoint = \"../model/sam_vit_h_4b8939.pth\"\n",
    "sam_model_type = \"vit_h\"\n",
    "sam_device = \"cuda\"\n",
    "\n",
    "image = cv2.imread('../test_images/4.jpg')\n",
    "image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "\n",
    "# 创建SAM\n",
    "sam = sam_model_registry[sam_model_type](checkpoint=sam_checkpoint)\n",
    "sam.to(device=sam_device)\n",
    "predictor = SamPredictor(sam)\n",
    "predictor.set_image(image)\n",
    "\n",
    "masks, scores, _ = predictor.predict()\n",
    "\n",
    "plt.figure(figsize=(10,10))\n",
    "plt.imshow(image)\n",
    "plt.axis('on')\n",
    "plt.show()  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 尝试Segament Automatically\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import cv2\n",
    "from segment_anything import SamAutomaticMaskGenerator, sam_model_registry\n",
    "import os\n",
    "os.environ[\"KMP_DUPLICATE_LIB_OK\"]=\"TRUE\"   # 一个粗暴的方法，解决库重复初始化问题\n",
    "\n",
    "def show_anns(anns):\n",
    "    if len(anns) == 0:\n",
    "        return\n",
    "    sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)\n",
    "    ax = plt.gca()\n",
    "    ax.set_autoscale_on(False)\n",
    "\n",
    "    img = np.ones((sorted_anns[0]['segmentation'].shape[0], sorted_anns[0]['segmentation'].shape[1], 4))\n",
    "    img[:,:,3] = 0\n",
    "    for ann in sorted_anns:\n",
    "        m = ann['segmentation']\n",
    "        color_mask = np.concatenate([np.random.random(3), [0.35]])\n",
    "        img[m] = color_mask\n",
    "    ax.imshow(img)\n",
    "\n",
    "\n",
    "sam_checkpoint = \"../model/sam_vit_h_4b8939.pth\"\n",
    "sam_model_type = \"vit_h\"\n",
    "sam_device = \"cuda\"\n",
    "\n",
    "image = cv2.imread('../dog.jpg')\n",
    "image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n",
    "\n",
    "        \n",
    "# 创建SAM\n",
    "sam = sam_model_registry[sam_model_type](checkpoint=sam_checkpoint)\n",
    "sam.to(device=sam_device)\n",
    "mask_generator = SamAutomaticMaskGenerator(\n",
    "    model=sam,\n",
    "    points_per_side=16,\n",
    "    points_per_batch=1,\n",
    "    pred_iou_thresh=0.86,\n",
    "    stability_score_thresh=0.92,\n",
    "    crop_n_layers=1,\n",
    "    crop_n_points_downscale_factor=2,\n",
    "    min_mask_region_area=100,  # Requires open-cv to run post-processing\n",
    ")\n",
    "\n",
    "masks = mask_generator.generate(image)\n",
    "\n",
    "\n",
    "show_anns(masks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "def show_mask(mask, ax, random_color=False):\n",
    "    if random_color:\n",
    "        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)\n",
    "    else:\n",
    "        color = np.array([30/255, 144/255, 255/255, 0.6])\n",
    "    h, w = mask.shape[-2:]\n",
    "    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)\n",
    "    ax.imshow(mask_image)\n",
    "    \n",
    "def show_points(coords, labels, ax, marker_size=375):\n",
    "    pos_points = coords[labels==1]\n",
    "    neg_points = coords[labels==0]\n",
    "    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)\n",
    "    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)   \n",
    "    \n",
    "def show_box(box, ax):\n",
    "    x0, y0 = box[0], box[1]\n",
    "    w, h = box[2] - box[0], box[3] - box[1]\n",
    "    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0,0,0,0), lw=2))  \n",
    "plt.figure(figsize=(10,10))\n",
    "\n",
    "for i, (mask, score) in enumerate(zip(masks, scores)):\n",
    "    plt.figure(figsize=(10,10))\n",
    "    plt.imshow(image)\n",
    "    show_mask(mask, plt.gca())\n",
    "    # show_points(input_point, input_label, plt.gca())\n",
    "    plt.title(f\"Mask {i+1}, Score: {score:.3f}\", fontsize=18)\n",
    "    plt.axis('off')\n",
    "    plt.show() "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
